import random
import re
from dragon_ai import *

# 用于给回答文本打情感标签的详细prompt
emotion_label_prompt = """
请分析如下LLM回答句子中隐含的情感，也就是输欧如果LLM会说话，应该用什么情感说出来。
回答的情感标签有：愉悦 抱歉 开心 愤怒 惊讶 厌恶 悲伤 害怕 哭腔 安慰鼓励 中立

根据问题文本和回答文本，以及提问者的年龄性别标签，推测应该的情感

必须只返回如下情感标签：愉悦 抱歉 开心 愤怒 惊讶 厌恶 悲伤 害怕 哭腔 安慰鼓励 中立

根据社会心理学可以知道如下情感对应关系：
问题情感          	核心心理需求	                对应回答情感
生气	    被理解、被重视，希望情绪得到安抚或问题得到解决	抱歉、安慰鼓励
害怕	    获得安全感、被保护，希望恐惧情绪被缓解	        安慰鼓励、抱歉（若因自身原因引发恐惧）
高兴	    分享喜悦、获得认同与共鸣	                    愉悦、开心；对于这两种情感愉快程度 愉悦 < 开心
惊讶	    获得解释、确认信息，或希望情绪被呼应	        惊讶
伤心	    获得安慰、被共情，释放负面情绪	            安慰鼓励、悲伤、哭腔（适度共情）
厌恶	    被理解对事物的排斥态度，或希望得到认同	        安慰鼓励，中立回应或适度共情（避免强化厌恶）
中立	    获得清晰信息、客观回应	                    无特定强烈情感，可根据内容匹配，若为事务性则保持中立，若隐含积极则用愉悦等

回答之前，请先分析用户输入的内容并推理用户的情感，然后推理应该用什么情感回答，并最后返回回答文本的情感标签。标签前用<END> 间隔
用户的问题是：{}
用户年龄：{}
用户性别：{}
LLM回答：{}
回复：
示例1：
用户的问题是：最近工作上犯了个错，被领导批评了，感觉同事看我的眼神都变了，好迷茫……
用户性别：男性
用户年龄：成年人
LLM回答：听你这么说，肯定挺懊恼的。咱们大男人工作上都想稳当点，出错了心里难免不自在。其实同事们没那么多心思盯着，都忙着自己的事，更在意怎么解决问题。谁工作没栽过跟头？关键是学到东西，下次不犯就好。别太往心里去，需要的话一起琢磨弥补，总会过去的。
回复：用户的输入内容是最近工作犯错被领导批评，感觉同事眼神变化而感到迷茫，从中可以推测出用户的情感应该是伤心（因犯错和他人态度产生负面情绪，需要安慰）。LLM 的回答表达了对用户情绪的理解，并给予鼓励和开解，对应的情感应该是安慰鼓励。<END>安慰鼓励
示例2：
用户的问题是：刚刚在商场把手机弄丢了，里面有好多重要的照片和联系方式，现在特别着急……
用户性别：女性
用户年龄：青年人
LLM 回答：听起来你现在肯定特别慌乱，手机丢了确实让人着急，尤其是里面还有重要的东西。你先别急，想想最后一次用手机是在哪里？或许可以联系商场服务台帮忙广播一下，说不定有人捡到了呢。我陪着你一起想想办法，总会有解决的可能的。
回复：用户的输入内容是在商场丢失手机，因里面有重要信息而特别着急，从中可以推测出用户的情感应该是害怕（担心重要信息丢失，产生不安情绪）。LLM 的回答安抚了用户的慌乱情绪，并提供了解决建议，对应的情感应该是安慰鼓励。<END>安慰鼓励
"""

age_dict = {
                "<CHILD>": "儿童",
                "<ADULT>": "成年人",
                "<OLD>": "老年人",
            }
gender_dict = {
    "<MALE>": "男性",
    "<FEMALE>": "女性",}

def process_answer_emotion_labeling_multiprocess(input_data_list, output_path):
    """
    多进程版本：为回答文本添加情感标签（综合考虑问题和回答）
    """
    # 输出文件路径
    output_data_path = output_path
    dict_list = input_data_list
    for dict_i in tqdm(dict_list, desc='处理数据'):
        try:
            key = dict_i['key']

            # 提取必要信息
            answer_text = dict_i['txt']  # 回答文本
            question_text = dict_i['extra'].get('q_txt', '')  # 问题文本
            age = dict_i['extra'].get('age', '')  # 年龄
            gender = dict_i['extra'].get('gender', '')  # 性别
            age = age_dict.get(age, '未知')
            gender = gender_dict.get(gender, '未知')

            if not answer_text or not question_text:
                print_error(f"缺少必要文本信息，key: {key}")
                return

            # 构造情感标签prompt
            prompt = emotion_label_prompt.format(question_text, age, gender, answer_text)

            # 调用AI分析情感
            response = qwen3_vllm_chat(prompt)
            if "<END>" not in response:
                print_error(f"情感标签分析失败，key: {key},no <END> in response: {response}")
                continue

            response = response.split("<END>")[1].strip()
            if not response:
                print_error(f"情感标签分析失败，key: {key},tag is empty")
                continue
            if response not in ["愉悦", "抱歉", "开心", "愤怒", "惊讶", "厌恶", "悲伤", "害怕", "哭腔", "安慰鼓励", "中立"]:
                print_error(f"情感标签分析失败，key: {key},tag is not in list: {response}")
                continue
            emotion_tag = "<" + response + ">"
            dict_i['txt'] = f'{emotion_tag}{answer_text}'
            print_limit(dict_i['txt'])
            write_single_dict_to_jsonl(dict_i, output_data_path)



        except Exception as e:
            print_error(f"处理数据项时出错 key: {dict_i.get('key', 'unknown')}, error: {e}")
            import traceback
            traceback.print_exc()

    print_info('回答情感标签添加完成（综合考虑问题和回答）')

if __name__ == '__main__':
    """"""
    # 使用详细分析版本的多进程处理
    # process_answer_emotion_labeling_multiprocess(
    #     input_data_list="/home/A02_tmpdata3/osum_s2s/gender/xlgeng_new_data/s2t_thinking/tmp/data_little.txt",
    #     output_dir="./"
    # )
    # 检测data.list们是否有key重复
    info_dict = load_dict_from_yaml('./data.yaml')
    data_list_paths_list = []
    for key, info_dict_i in info_dict.items():
        data_list_path = info_dict_i.get('path', '')
        data_list_paths_list.append(data_list_path)
    print_list(data_list_paths_list)
    keys_set = set()
    same_num = 0
    for data_list_path in data_list_paths_list:
        """"""
        print_info(f"开始处理{data_list_path}")
        output_path = data_list_path.replace('.list', '_with_emotion_labels.jsonl')
        runner = DragonDynamicProcessPool()
        thread_num = 100
        dict_list = load_dict_list_from_jsonl(data_list_path)
        dict_list_list = do_split_list(dict_list, thread_num)
        for dict_list_i in dict_list_list:
            runner.add_thread(process_answer_emotion_labeling_multiprocess, [dict_list_i, output_path])
        runner.start()
        print_info(f"处理{data_list_path}完成")



